#!/usr/bin/env python33
#coding=utf8
#

import threading
import socket
import re
import sys
import os
import requests

import models

class Scan(object):
    def is_alive(self,host,port,timeout=5):
        print(" [+] %.15s %.5s Scan..."%(host,port),end='\r')
        """检测端口是否开放"""
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        try:
            sock.settimeout(timeout)
            sock.connect((host,int(port)))
            sock.close()
            return True
        except:
            return False

    def start(self):
        if self.hosts and self.ports:
            is_conver = True
            if len(self.hosts)>len(self.ports): #扫描多个主机的同一端口
                is_conver = False
                iterlist = self.hosts
                tasklist = self.ports
            else:                               #扫描同一主机的多个端口
                iterlist = self.ports
                tasklist = self.hosts
            for host in tasklist:
                iters = iter(iterlist)
                loop = True
                while loop:
                    threads=[]
                    for _ in range(self.thread_num):
                        try:
                            port = iters.__next__()
                        except StopIteration:
                            loop = False
                            break
                        p,h = host,port
                        if is_conver:
                           h,p = host,port
                        t = threading.Thread(target=self.save,args=(h,p,))
                        threads.append(t)
                    loops = range(len(threads))
                    for i in loops:
                        threads[i].start()
                    for i in loops:
                        threads[i].join()

class ScanTask(Scan):
    def __init__(self):
        self.hosts = []         #IP列表
        self.ports = []         #端口列表
        self.is_forced = 1      #是否强制扫描
        self.is_http = 1        #是否获取banner
        self.thread_num = 10    #线程数
        self.sock_timeout = 5   #socket扫描超时时间
        self.http_timeout = 10  #http扫描超时时间

    def setinit(self,*args,**kwargs):
        for key,value in kwargs.items():
            setattr(self,key,value)

    def save(self,host,port):
        if self.is_alive(host,port,self.sock_timeout):
            print(" [+] %s %s Is Alive"%(host,port))
            taskip,created = models.IPResult.create_or_get(ip=host)
            result,created = models.PortResult.get_or_create(host=taskip,port=port)
            if self.is_http:
                try:
                    res = requests.get("HTTP://%s:%s"%(host,port))
                    result.status_code = res.status_code
                    result.headers = res.headers
                    result.content = res.text
                    result.banner = res.headers.get('Server','')
                    result.title = re.findall("<title>(.*)</title>",res.text)[0]
                except Exception as e:
                    result.title = str(e)
            result.save()

def main(taskid):
    from util import get_hosts,get_ports
    for task in models.IPTask.select().join(
                models.Task).where(
                models.Task.task_id==taskid):#.get()
        ST=ScanTask()
        ST.setinit(
            hosts = get_hosts(task.task_host),
            ports = get_ports(task.task_port),
            is_http = task.task_ishttp,
            is_forced = task.task_isforced)
        ST.start()

if __name__=="__main__":
    main(1)



